
#include <malloc.h>
#include <math.h>
#include <stdio.h>

#define TOLERATE 1e-3

void CG( double** A, double* x, double* b, int scale );
void JacobiPreconditionedCG( double** A, double* x, double* b, int scale );
void Jacobi( double** A, double* x, double* b, int scale );
void GaussSeidel( double** A, double* x, double* b, int scale );
void RedBlackGaussSeidel( double** A, double* x, double* b, int nx, int ny );
void DiscretizePDE( double** A, int scale_x, int scale_y, int depth, double hx, double hy );
float ComputeResidual( double** A, double* x, double* b, int scale );

void main(int argc, char *argv[])
{
	int nx, ny;

	nx = 60;
	ny = 60;

	int scale = nx * ny;

	double** A = (double**)malloc(scale*sizeof(double*));
	double* x = (double*)malloc(scale*sizeof(double));
	double* b = (double*)malloc(scale*sizeof(double));

	for (int j=0;j<scale;j++)
	{
		A[j] = (double*)malloc(scale*sizeof(double));
		x[j] = 0;
		b[j] = j;

		for (int k=0;k<scale;k++)
		{
			A[j][k] = 0;
		}
	}

	DiscretizePDE(A,nx,ny,0,1,1);

	JacobiPreconditionedCG(A,x,b,scale);
	//RedBlackGaussSeidel(A,x,b,nx,ny);

}

void CG( double** A, double* x_, double* b, int scale )
{
	double* residual = (double*)malloc(scale*sizeof(double));
	double* residual_cache = (double*)malloc(scale*sizeof(double));
	double* pbuffer = (double*)malloc(scale*sizeof(double));
	double* buffer = (double*)malloc(scale*sizeof(double));
	double* x = (double*)malloc(scale*sizeof(double));

	for (int i=0;i<scale;i++)
	{
		double sum=0;

		for (int j=0;j<scale;j++)
		{
			sum+=A[i][j]*x[j];
		}

		residual[i]=b[i]-sum;
		pbuffer[i]=residual[i];
	}

	int loop;

	for (loop=0;loop<1000;loop++)
	{
		double alpha, sumU=0, sumD=0;

		for (int i=0;i<scale;i++)
		{
			sumU+=residual[i]*residual[i];
			double sumT=0;

			for (int j=0;j<scale;j++)
			{
				sumT+=pbuffer[j]*A[j][i];
			}

			sumD+=sumT*pbuffer[i];
		}

		alpha=sumU/sumD;

		for (int i=0;i<scale;i++)
		{
			buffer[i]=x[i]+alpha*pbuffer[i];

			double sum=0;

			for (int j=0;j<scale;j++)
			{
				sum+=A[i][j]*pbuffer[j];
			}

			residual_cache[i]=residual[i]-alpha*sum;
		}

		/************************************************************************/
		/* check here                                                                     */
		/************************************************************************/

		double r = ComputeResidual(A,x,b,scale);

		if (r<TOLERATE)
		{
			double* swap;

			swap = x;
			x = buffer;
			buffer = x;

			for (int m=0;m<scale;m++)
			{
				x_[m] = x[m];
			}

			free(residual);
			free(residual_cache);
			free(pbuffer);
			free(buffer);
			free(x);

			printf("%d\n",loop);

			return;
		}else{
			printf("%f\n",r);
		}

		double beta;
		sumU=0;
		sumD=0;

		for (int i=0;i<scale;i++)
		{
			sumU+=residual_cache[i]*residual_cache[i];
			sumD+=residual[i]*residual[i];
		}

		beta=sumU/sumD;

		for (int i=0;i<scale;i++)
		{
			pbuffer[i]=residual_cache[i]+beta*pbuffer[i];
		}

		double* swap;

		swap = residual;
		residual = residual_cache;
		residual_cache = swap;

		swap = x;
		x = buffer;
		buffer = swap;
	}

	free(residual);
	free(residual_cache);
	free(pbuffer);
	free(buffer);
	free(x);
}

void DiscretizePDE( double** A, int scale_x, int scale_y, int depth, double hx, double hy )
{
	int index, current;

	for (int y=0;y<scale_y;y++)
	{
		for (int x=0;x<scale_x;x++)
		{
			current = y * scale_x + x;

			A[current][current] = 4;

			if (x-1>=0)
			{
				index = y * scale_x + x - 1;

				//A[current][current] ++;
				A[current][index] = -1.0 / (hx * hx);
			}

			if (y-1>=0)
			{
				index = (y - 1) * scale_x + x;

				//A[current][current] ++;
				A[current][index] = -1.0 / (hy * hy);
			}

			if (x+1<scale_x)
			{
				index = y * scale_x + x + 1;

				//A[current][current] ++;
				A[current][index] = -1.0 / (hx * hx);
			}

			if (y+1<scale_y)
			{
				index = (y + 1) * scale_x + x;

				//A[current][current] ++;
				A[current][index] = -1.0 / (hy * hy);
			}

			A[current][current] /= (hx * hy);
		}
	}
}

void Jacobi( double** A, double* x, double* b, int scale )
{
	double* buffer = (double*)malloc(scale*sizeof(double));

	for (int loop=0;loop<10000;loop++)
	{
		for (int i=0;i<scale;i++)
		{
			double sum = 0;

			for (int j=0;j<scale;j++)
			{
				if (i!=j)
				{
					sum += A[i][j] * x[j];
				}
			}

			if (A[i][i]!=0)
			{
				buffer[i] = (b[i] - sum) / A[i][i];
			}
		}

		for (int i=0;i<scale;i++)
		{
			x[i]=buffer[i];
		}

		double r = ComputeResidual(A,x,b,scale);

		if (r<TOLERATE)
		{
			printf("%d\n",loop);
			break;
		}else{
			printf("%f\n",r);
		}

	}

	free(buffer);
}

void GaussSeidel( double** A, double* x, double* b, int scale )
{
	for (int loop=0;loop<10000;loop++)
	{
		double residual = 0;

		for (int i=0;i<scale;i++)
		{
			double sum = 0;

			for (int j=0;j<scale;j++)
			{
				if (i!=j)
				{
					sum += A[i][j] * x[j];
				}
			}

			if (A[i][i]!=0)
			{
				x[i] = (b[i] - sum) / A[i][i];
			}
		}

		residual = ComputeResidual(A,x,b,scale);

		if (residual<TOLERATE)
		{
			printf("%d\n",loop);
			break;
		}else{
			printf("%f\n",residual);
		}
	}
}

float ComputeResidual( double** A, double* x, double* b, int scale )
{
	double r = 0;

	for (int i=0;i<scale;i++)
	{
		double sum = 0;

		for (int j=0;j<scale;j++)
		{
			sum += A[i][j] * x[j];
		}
		
		r += (sum - b[i]) * (sum - b[i]);
	}

	r = sqrt(r);

	return r;
}

void RedBlackGaussSeidel( double** A, double* x, double* b, int nx, int ny )
{
	int scale = nx * ny;

	for (int loop=0;loop<10000;loop++)
	{
		double residual = 0;

		for (int sy=0;sy<ny;sy++)
		{
			for (int sx=0;sx<nx;sx++)
			{
				if ((sx+sy)%2==0)
				{
					int index = sy * nx + sx;

					double sum = 0;

					for (int k=0;k<scale;k++)
					{
						if (k!=index)
						{
							sum += A[index][k] * x[k];
						}
					}

					if (A[index][index]!=0)
					{
						x[index] = (b[index] - sum) / A[index][index];
					}
				}
			}
		}

		for (int sy=0;sy<ny;sy++)
		{
			for (int sx=0;sx<nx;sx++)
			{
				if ((sx+sy)%2==1)
				{
					int index = sy * nx + sx;

					double sum = 0;

					for (int k=0;k<scale;k++)
					{
						if (k!=index)
						{
							sum += A[index][k] * x[k];
						}
					}

					if (A[index][index]!=0)
					{
						x[index] = (b[index] - sum) / A[index][index];
					}
				}
			}
		}

		/*for (int i=0;i<scale;i++)
		{
			double sum = 0;

			for (int j=0;j<scale;j++)
			{
				if (i!=j)
				{
					sum += A[i][j] * x[j];
				}
			}

			if (A[i][i]!=0)
			{
				x[i] = (b[i] - sum) / A[i][i];
			}
		}*/

		residual = ComputeResidual(A,x,b,scale);

		if (residual<TOLERATE)
		{
			printf("%d\n",loop);
			break;
		}else{
			printf("%f\n",residual);
		}
	}
}

void JacobiPreconditionedCG( double** PA, double* x_, double* Pb, int scale )
{
	double** A = (double**)malloc(scale*sizeof(double*));
	double* b = (double*)malloc(scale*sizeof(double));

	for (int j=0;j<scale;j++)
	{
		A[j] = (double*)malloc(scale*sizeof(double));
		b[j] = Pb[j] / PA[j][j];

		for (int k=0;k<scale;k++)
		{
			A[j][k] = PA[j][k] / PA[j][j];
		}
	}

	double* residual = (double*)malloc(scale*sizeof(double));
	double* residual_cache = (double*)malloc(scale*sizeof(double));
	double* pbuffer = (double*)malloc(scale*sizeof(double));
	double* buffer = (double*)malloc(scale*sizeof(double));
	double* x = (double*)malloc(scale*sizeof(double));

	for (int i=0;i<scale;i++)
	{
		double sum=0;

		for (int j=0;j<scale;j++)
		{
			sum+=A[i][j]*x[j];
		}

		residual[i]=b[i]-sum;
		pbuffer[i]=residual[i];
	}

	int loop;

	for (loop=0;loop<1000;loop++)
	{
		double alpha, sumU=0, sumD=0;

		for (int i=0;i<scale;i++)
		{
			sumU+=residual[i]*residual[i];
			double sumT=0;

			for (int j=0;j<scale;j++)
			{
				sumT+=pbuffer[j]*A[j][i];
			}

			sumD+=sumT*pbuffer[i];
		}

		alpha=sumU/sumD;

		for (int i=0;i<scale;i++)
		{
			buffer[i]=x[i]+alpha*pbuffer[i];

			double sum=0;

			for (int j=0;j<scale;j++)
			{
				sum+=A[i][j]*pbuffer[j];
			}

			residual_cache[i]=residual[i]-alpha*sum;
		}

		/************************************************************************/
		/* check here                                                                     */
		/************************************************************************/

		double r = ComputeResidual(PA,x,Pb,scale);

		if (r<TOLERATE)
		{
			double* swap;

			swap = x;
			x = buffer;
			buffer = x;

			for (int m=0;m<scale;m++)
			{
				x_[m] = x[m];
			}

			free(residual);
			free(residual_cache);
			free(pbuffer);
			free(buffer);
			free(x);

			printf("%d\n",loop);

			return;
		}else{
			printf("%f\n",r);
		}

		double beta;
		sumU=0;
		sumD=0;

		for (int i=0;i<scale;i++)
		{
			sumU+=residual_cache[i]*residual_cache[i];
			sumD+=residual[i]*residual[i];
		}

		beta=sumU/sumD;

		for (int i=0;i<scale;i++)
		{
			pbuffer[i]=residual_cache[i]+beta*pbuffer[i];
		}

		double* swap;

		swap = residual;
		residual = residual_cache;
		residual_cache = swap;

		swap = x;
		x = buffer;
		buffer = swap;
	}

	free(residual);
	free(residual_cache);
	free(pbuffer);
	free(buffer);
	free(x);

}
